import matplotlib.pyplot as plt

def plot_images(tensor_list, batch_index=0, use_max_variance=False, channel_index=0):
    num_images = len(tensor_list)
    fig, axes = plt.subplots(1, num_images, figsize=(15, 5))
    if num_images == 1:
        axes = [axes]

    for ax, tensor in zip(axes, tensor_list):
        if tensor.dim() == 4:
            if use_max_variance:
                variances = tensor[batch_index].var(dim=(1, 2))
                channel_index = variances.argmax().item()
            image = tensor[batch_index, channel_index].detach().cpu().numpy()
        elif tensor.dim() == 3:
            if use_max_variance:
                variances = tensor.var(dim=(1, 2))
                channel_index = variances.argmax().item()
            image = tensor[channel_index].cpu().numpy()
        else:
            raise ValueError("Unsupported tensor shape.")

        if image.max() > 1:
            image = image / 255.0

        ax.imshow(image, cmap='gray' if tensor.dim() == 3 else None)
        ax.axis('off')

    plt.tight_layout()
    plt.show()

def plot_tensor(tensor_3d):
    C, H, W = tensor_3d.shape
    
    for i in range(C):
        plt.imshow(tensor_3d[i], cmap='viridis', interpolation='nearest')
        plt.colorbar()
        plt.title(f'Channel {i+1} of Tensor')
        plt.show()
